package online.chazi.bootplus.kafka.consumer.handler;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import online.chazi.bootplus.commons.utils.SleepUtils;
import online.chazi.bootplus.kafka.consumer.KafkaConsumerOffset;
import online.chazi.bootplus.kafka.consumer.config.properties.KafkaConsumerProperties;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;

import javax.annotation.Resource;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

/**
 * 说明：作者很懒，什么都没留下
 * Created by 叉子同学 on 2020-04-30 13:01
 */
@Slf4j
public class KafkaConsumerPoller<K, V> {

    private boolean interrupt;
    private KafkaConsumerHandler<K, V> handler;
    private Collection<String> topics;
    private Thread thread;

    @Resource
    private KafkaConsumerProperties kafkaConsumerProperties;
    @Resource
    private KafkaConsumer<K, V> kafkaConsumer;

    public void stop() {
        interrupt = true;
        try {
            thread.join();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void runWithHandler(KafkaConsumerHandler<K, V> handler, Collection<String> topics) {
        this.handler = handler;
        this.topics = topics;
        thread = new Thread(new Runner());
        thread.setDaemon(false);
        thread.start();
    }

    private class Runner implements Runnable {

        private String consumerId = kafkaConsumerProperties.getGroupId();

        private void init() {
            kafkaConsumer.subscribe(
                    topics,
                    new ConsumerRebalanceListener() {
                        @Override
                        public void onPartitionsRevoked(Collection<TopicPartition> collection) {
                            //所有partition一起更新入库
                            List<KafkaConsumerOffset> offsetList = new ArrayList<>();
                            collection.forEach(partition -> {
                                KafkaConsumerOffset offset = new KafkaConsumerOffset();
                                offset.setTopic(partition.topic());
                                offset.setPartition(partition.partition());
                                offset.setPosition(kafkaConsumer.position(partition));
                                offset.setConsumerId(consumerId);
                                //
                                offsetList.add(offset);
                            });
                            if (offsetList.size() > 0) {
                                handler.rebalanceOffset(offsetList);
                            }
                        }

                        @Override
                        public void onPartitionsAssigned(Collection<TopicPartition> collection) {
                            seek(collection);
                        }
                    });
        }

        private void seek(Collection<TopicPartition> collection) {
            collection.forEach(partition -> {
                KafkaConsumerOffset offset = handler.seekOffset(
                        consumerId,
                        partition.topic(),
                        partition.partition());
                if (offset != null) {
                    if (offset.getPosition() == null) {
                        log.debug("定位Kafka位置offset, seekToEnd, topic={},partition={}", partition.topic(), partition.partition());
                        kafkaConsumer.seekToEnd(Collections.singleton(partition));
                    } else {
                        log.debug("定位Kafka位置offset, seek, offset={}", JSON.toJSONString(offset));
                        kafkaConsumer.seek(partition, offset.getPosition());
                    }
                }
            });
        }

        @Override
        public void run() {
            log.info("消费任务开启了。");
            init();
            while (!interrupt) try {
                if (handler.goAhead()) {
                    consume();
                }
                SleepUtils.sleep(handler.sleep());
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
            log.info("优雅的结束。");
        }

        private void consume() {

            List<KafkaConsumerOffset> offsets = new ArrayList<>();
            ConsumerRecords<K, V> records = kafkaConsumer.poll(Duration.ofMillis(kafkaConsumerProperties.getFetchMaxWait()));

            if (records.isEmpty()) handler.recordsEmpty();
            else records.partitions().forEach(partition -> {
                List<ConsumerRecord<K, V>> partitionRecords = records.records(partition);
                handler.consumeCount(partition.topic(), partition.partition(), partitionRecords.size());
                partitionRecords.forEach(record -> {
                    try {

                        V value = record.value();
                        handler.consume(record.topic(), record.partition(), record.offset(), value);

                    } catch (Exception e) {
                        log.error("处理异常：" + e.getMessage(), e);
                        handler.caughtException(e);
                    }
                });

                long lastOffsetPosition = partitionRecords.get(partitionRecords.size() - 1).offset() + 1;
                if (kafkaConsumerProperties.getEnableAutoCommit()) {
                    kafkaConsumer.commitSync(Collections.singletonMap(partition, new OffsetAndMetadata(lastOffsetPosition)));
                }

                KafkaConsumerOffset offset = new KafkaConsumerOffset();
                offset.setConsumerId(consumerId);
                offset.setTopic(partition.topic());
                offset.setPartition(partition.partition());
                offset.setPosition(lastOffsetPosition);
                offsets.add(offset);
            });

            if (offsets.size() > 0) handler.rebalanceOffset(offsets);

        }

    }

}
